from datetime import datetime, timedelta
import numpy as np

class PeriodPredictor:
    def __init__(self, user):
        self.user = user
        
    def predict_next_period(self):
        if not self.user.period_records:
            return None
            
        # 获取最近的经期记录
        dates = [datetime.strptime(record['start_date'], '%Y-%m-%d') 
                for record in self.user.period_records]
        
        if len(dates) < 2:
            # 如果只有一条记录，使用用户提供的平均周期长度
            last_date = dates[-1]
            predicted_date = last_date + timedelta(days=self.user.avg_cycle_length)
        else:
            # 计算实际周期长度
            intervals = [(dates[i+1] - dates[i]).days 
                        for i in range(len(dates)-1)]
            
            # 使用最近三次周期的加权平均
            weights = [0.5, 0.3, 0.2]
            if len(intervals) >= 3:
                avg_interval = sum(w * i for w, i in zip(weights, intervals[-3:]))
            else:
                avg_interval = sum(intervals) / len(intervals)
                
            predicted_date = dates[-1] + timedelta(days=round(avg_interval))
            
        return predicted_date
        
    def get_reminder_dates(self):
        next_period = self.predict_next_period()
        if not next_period:
            return []
            
        return [next_period - timedelta(days=i) for i in range(1, 4)] 